import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 1. 创造数据，数据集
class1_points = np.array(
    [[6.5, 4.3], [4.5, 6.4], [1.3, 5.1], [1.7, 4.4], [4.8, 5.7], [5.4, 5.6], [1.8, 4.9], [1.2, 3.8],
     [2.8, 5.7], [6.4, 3.8],
     [4.5, 5.9], [5.3, 6.0], [5.9, 5.0], [1.7, 4.6], [2.3, 5.7], [3.4, 6.1], [5.9, 4.4], [5.4, 5.1],
     [5.2, 5.2], [5.6, 5.4],
     [4.2, 6.2], [1.4, 3.7], [3.6, 6.3], [4.8, 6.0], [4.8, 6.0], [5.0, 6.1], [5.8, 5.1], [1.6, 4.5],
     [1.5, 5.1], [2.2, 6.0],
     [5.1, 5.8], [3.8, 6.3], [2.0, 5.7], [2.1, 5.6], [2.0, 5.1], [1.0, 4.9], [3.0, 6.3], [6.0, 4.2],
     [2.3, 6.3], [4.8, 6.1],
     [1.8, 5.1], [2.2, 5.7], [6.3, 4.3], [5.7, 5.3], [5.6, 5.5], [3.0, 6.1], [6.1, 3.7], [6.3, 4.7],
     [3.4, 6.1], [5.2, 5.7],
     [5.8, 3.7], [0.7, 4.6], [4.9, 6.2], [1.8, 5.1], [4.6, 5.9], [1.5, 5.0], [1.4, 4.4], [4.0, 6.4],
     [5.3, 5.8], [4.6, 6.1],
     [3.5, 6.0], [6.2, 4.6], [4.5, 6.0], [2.6, 6.1], [5.9, 5.0], [2.8, 6.4], [2.4, 6.0], [5.3, 6.0],
     [2.0, 5.7], [1.2, 3.7],
     [2.8, 5.9], [2.5, 5.5], [6.3, 4.6], [1.2, 3.7], [6.3, 4.4], [6.0, 4.8], [1.5, 4.2], [6.4, 4.2],
     [1.3, 4.6], [2.0, 5.2],
     [1.9, 5.2], [1.6, 5.4], [5.5, 5.7], [3.5, 6.6], [1.7, 5.0], [6.2, 4.6], [6.1, 4.5], [4.1, 5.9],
     [6.1, 4.9], [1.7, 5.2],
     [3.5, 6.2], [2.9, 6.4], [5.0, 5.8], [2.5, 5.8], [3.1, 6.0], [2.0, 5.1], [2.6, 5.7], [6.1, 4.0],
     [6.5, 4.4], [5.4, 6.1],
     [5.9, 4.1], [4.7, 5.9], [2.4, 6.5], [4.5, 6.4], [5.9, 4.6], [0.9, 3.9], [3.6, 6.3], [3.7, 6.3],
     [1.6, 4.3], [6.0, 5.7],
     [4.2, 6.3], [1.8, 5.2], [2.7, 5.9], [2.4, 5.5], [6.4, 3.8], [5.2, 6.1], [6.2, 4.7], [4.2, 6.5],
     [5.7, 3.6], [3.9, 6.1],
     [1.1, 4.6], [5.5, 5.3], [2.0, 5.9], [5.2, 5.4], [5.7, 5.2], [5.3, 5.0], [1.4, 4.1], [2.8, 6.6],
     [3.6, 6.3], [1.1, 4.3],
     [5.5, 5.2], [3.9, 6.9], [6.2, 4.2], [5.5, 5.5], [1.6, 4.1], [1.1, 3.9], [1.4, 4.9], [4.5, 6.1],
     [1.7, 5.0], [1.9, 4.7],
     [5.8, 5.7], [4.8, 5.6], [3.2, 5.7], [6.3, 4.0], [1.6, 4.2], [1.8, 5.1], [1.9, 5.5], [2.9, 5.6],
     [1.0, 3.8], [5.9, 5.5],
     [2.6, 5.6], [5.3, 5.4], [1.5, 5.0], [3.2, 6.1], [1.0, 4.1], [1.9, 5.8], [3.3, 6.2], [6.1, 3.9],
     [2.9, 5.8], [4.8, 5.9],
     [6.0, 4.4], [3.6, 6.2], [1.6, 5.1], [5.6, 5.0], [4.0, 6.2], [6.2, 4.3], [4.2, 6.4], [4.0, 6.1],
     [5.5, 5.1], [4.3, 6.1],
     [4.5, 5.8], [3.7, 6.7], [1.6, 5.6], [5.7, 4.6], [1.6, 4.9], [6.2, 5.7], [2.8, 6.2], [2.1, 5.7],
     [5.8, 6.2], [1.5, 5.0],
     [5.6, 5.6], [4.1, 5.7], [1.8, 4.6], [6.4, 4.1], [1.2, 3.8], [2.4, 6.0], [1.5, 5.2], [6.0, 3.9],
     [5.9, 4.7], [1.9, 5.5],
     [2.3, 5.5], [6.1, 4.4], [2.0, 5.2], [1.8, 5.5], [4.6, 6.3], [3.4, 6.2], [4.7, 6.3], [3.1, 6.1],
     [3.8, 6.3], [5.7, 5.5],
     [1.9, 5.4], [4.7, 5.9], [6.0, 4.2], [4.5, 6.5], [1.3, 4.2], [5.1, 6.0], [1.8, 5.2], [4.0, 6.4],
     [5.8, 5.6], [1.2, 3.9],
     [6.1, 5.4], [1.7, 4.9], [6.3, 5.0], [5.2, 5.0], [3.0, 6.4], [1.6, 4.8], [1.5, 5.2], [4.7, 6.3],
     [1.5, 4.8], [5.3, 5.8],
     [4.3, 5.9], [3.2, 6.3], [2.4, 5.5], [2.6, 5.4], [1.2, 3.9], [4.8, 6.3], [6.2, 4.6], [1.3, 5.3],
     [6.6, 4.1], [2.9, 6.3],
     [3.3, 6.1], [6.0, 5.3], [1.5, 4.9], [5.6, 5.7], [5.9, 4.5], [4.9, 6.1], [6.0, 4.6], [5.0, 5.4],
     [3.4, 6.1], [5.9, 4.9],
     [2.8, 5.4], [1.9, 5.3], [3.2, 5.8], [1.2, 4.7], [3.1, 6.3], [1.2, 4.0], [6.0, 5.7], [2.7, 6.0],
     [3.4, 6.0], [5.9, 5.4]])
class2_points = np.array(
    [[6.5, 2.5], [6.4, 2.3], [6.6, 2.8], [7.0, 2.6], [4.3, 2.9], [4.1, 3.7], [3.9, 3.3], [7.2, 2.7],
     [3.8, 4.5], [4.0, 4.7],
     [4.0, 3.9], [8.3, 3.8], [6.5, 3.1], [8.0, 3.6], [7.9, 3.4], [6.8, 2.5], [4.0, 4.4], [7.0, 2.6],
     [7.7, 3.1], [6.0, 2.1],
     [6.7, 2.7], [8.7, 4.2], [4.0, 3.9], [5.9, 2.2], [6.3, 2.7], [7.3, 2.9], [5.0, 2.6], [8.1, 3.9],
     [4.2, 4.0], [5.1, 2.5],
     [8.2, 3.3], [7.1, 2.9], [5.0, 3.0], [7.1, 2.3], [4.8, 3.1], [3.5, 4.4], [8.3, 3.3], [5.2, 3.0],
     [6.1, 2.2], [6.8, 2.2],
     [3.9, 4.9], [8.6, 3.6], [6.0, 2.3], [4.1, 4.0], [5.2, 2.8], [8.2, 3.5], [8.1, 3.4], [8.7, 4.9],
     [5.0, 2.4], [5.0, 2.6],
     [8.0, 3.0], [8.4, 4.3], [5.3, 2.7], [8.7, 5.1], [5.6, 2.5], [5.4, 2.7], [3.8, 4.5], [9.1, 4.3],
     [8.8, 4.1], [4.7, 3.3],
     [8.4, 4.6], [8.3, 4.5], [7.0, 2.7], [6.4, 2.3], [5.2, 2.5], [7.0, 2.2], [8.6, 3.3], [7.5, 3.0],
     [4.0, 3.9], [7.6, 3.0],
     [7.0, 2.7], [4.3, 3.1], [5.7, 2.8], [3.8, 4.3], [4.9, 3.1], [4.1, 3.3], [7.0, 2.3], [5.1, 2.9],
     [8.9, 4.5], [6.0, 2.7],
     [7.4, 2.6], [8.7, 4.7], [8.6, 4.5], [7.7, 3.0], [8.9, 5.0], [4.1, 4.0], [3.9, 4.8], [3.7, 3.8],
     [5.5, 2.3], [7.5, 3.4],
     [4.2, 3.3], [4.1, 3.5], [7.8, 3.1], [3.8, 4.7], [5.2, 3.3], [3.5, 4.7], [3.5, 4.8], [3.9, 4.2],
     [6.7, 3.1], [7.9, 3.0],
     [8.6, 4.1], [8.5, 4.4], [7.3, 2.6], [3.4, 4.7], [8.7, 3.9], [7.6, 3.0], [4.6, 3.1], [4.8, 2.7],
     [4.5, 2.5], [7.4, 2.9],
     [5.1, 2.7], [6.9, 2.7], [7.6, 2.6], [9.0, 5.0], [7.1, 2.2], [5.0, 2.7], [5.6, 2.4], [3.6, 4.8],
     [6.0, 2.4], [6.9, 2.9],
     [8.3, 4.9], [3.9, 4.0], [4.9, 3.1], [8.7, 3.9], [6.3, 2.4], [6.8, 2.5], [5.8, 2.1], [4.5, 4.1],
     [4.7, 3.2], [6.3, 2.6],
     [8.8, 4.8], [8.6, 4.1], [4.5, 3.8], [3.6, 4.3], [8.8, 5.0], [4.2, 3.9], [8.6, 4.4], [8.8, 4.0],
     [5.0, 3.4], [6.4, 2.5],
     [4.6, 2.6], [6.0, 2.6], [8.1, 3.5], [8.7, 4.5], [4.8, 2.8], [5.9, 2.7], [6.8, 2.6], [8.9, 4.6],
     [6.4, 2.6], [6.9, 2.5],
     [8.8, 3.3], [3.7, 4.0], [8.3, 4.0], [3.6, 4.3], [7.2, 2.2], [8.8, 4.4], [8.7, 4.7], [3.8, 4.4],
     [8.1, 3.4], [3.5, 4.7],
     [8.7, 4.1], [4.3, 3.8], [3.6, 4.0], [5.0, 2.7], [7.7, 3.2], [8.4, 3.2], [4.3, 3.7], [8.6, 4.3],
     [7.5, 3.2], [8.3, 3.8],
     [4.9, 2.9], [5.4, 2.4], [3.9, 4.9], [8.9, 3.6], [8.3, 3.4], [8.2, 3.3], [7.8, 2.8], [8.2, 3.2],
     [8.9, 4.8], [8.6, 3.8],
     [3.9, 5.3], [4.4, 4.6], [7.8, 3.0], [6.9, 2.7], [7.7, 3.0], [3.7, 3.7], [6.6, 3.0], [5.3, 2.6],
     [4.4, 4.1], [8.1, 3.6],
     [8.5, 3.4], [8.0, 3.7], [5.2, 2.7], [7.3, 2.8], [4.1, 4.0], [8.5, 3.6], [7.5, 2.4], [3.9, 3.8],
     [5.9, 2.5], [6.6, 2.9],
     [4.4, 3.4], [4.8, 3.3], [4.4, 3.1], [8.7, 4.8], [6.2, 2.7], [5.0, 3.2], [5.6, 2.7], [8.5, 4.2],
     [4.2, 3.5], [4.0, 3.1],
     [3.8, 4.1], [5.3, 2.2], [4.9, 3.3], [5.7, 3.1], [4.4, 3.5], [5.3, 2.8], [4.2, 3.3], [8.4, 3.6],
     [8.1, 3.5], [3.8, 4.4],
     [3.6, 4.3], [4.3, 4.6], [7.9, 3.1], [8.9, 4.9], [7.8, 3.2], [4.1, 3.7], [4.8, 3.1], [3.7, 4.3],
     [8.5, 3.8], [5.2, 2.7],
     [7.3, 2.8], [6.5, 2.6], [8.4, 4.3], [8.2, 4.0], [7.2, 2.9], [3.7, 4.2], [7.6, 2.6], [4.3, 4.7],
     [4.5, 3.5], [4.0, 4.2],
     [6.4, 2.7], [6.3, 2.6], [8.9, 3.9], [5.8, 2.3], [6.1, 2.6], [4.1, 3.7], [8.2, 3.1], [9.1, 4.5],
     [3.7, 4.1], [6.3, 2.7]])
# 将原始数据分成训练和测试两类
np.random.shuffle(class1_points)
np.random.shuffle(class2_points)
# 取前15%作为测试数据集
split_index_class1 = int(0.15 * len(class1_points))
split_index_class2 = int(0.15 * len(class2_points))

class1_train_points = np.array(class1_points[split_index_class1:])
class1_test_points = np.array(class1_points[:split_index_class1])
class2_train_points = np.array(class2_points[split_index_class2:])
class2_test_points = np.array(class2_points[:split_index_class2])

# 合并和创造训练的点和标签
train_point = np.concatenate((class1_train_points, class2_train_points))
train_label = np.concatenate((np.zeros(len(class1_train_points)), np.ones(len(class2_train_points))))

# 合并和创造测试的点和标签
test_point = np.concatenate((class1_test_points, class2_test_points))
test_label = np.concatenate((np.zeros(len(class1_test_points)), np.ones(len(class2_test_points))))

# 将Numpy数组转成tensor
point_train_tensor = torch.tensor(train_point, dtype=torch.float32)
label_train_tensor = torch.tensor(train_label, dtype=torch.long)
point_test_tensor = torch.tensor(test_point, dtype=torch.float32)
label_test_tensor = torch.tensor(test_label, dtype=torch.long)

# 2.定义前向模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer1 = nn.Linear(2, 8)
        self.layer2 = nn.Linear(8, 8)
        self.layer3 = nn.Linear(8, 2)

    def forward(self, x):
        x = torch.tanh(self.layer1(x))
        x = torch.tanh(self.layer2(x))
        x = torch.softmax(self.layer3(x), dim=1)
        return x

model = Model()

# 3.定义损失函数和优化器
cri = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 初始化绘图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
train_loss_list, test_loss_list, train_epoch_list, test_epoch_list = [], [], [], []
# 新建网格
xx, yy = np.meshgrid(np.linspace(0, 10, 100), np.linspace(0, 10, 100))
grid_point = np.c_[xx.ravel(), yy.ravel()]
grid_point = torch.tensor(grid_point, dtype=torch.float32)

# 4.开始迭代
epoches = 1000
for epoch in range(1, epoches+1):
    # 前向传播
    y_pred = model(point_train_tensor)
    loss = cri(y_pred, label_train_tensor)  # 训练的损失

    # 更新参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 5.显示频率设置
    if epoch == 1 or epoch % 20 == 0:
        # 6.绘图
        print(f"epoch: {epoch}, loss: {loss}")
        # 更新完才拿到test损失
        y_pred_test = model(point_test_tensor)
        test_loss = cri(y_pred_test, label_test_tensor)
        train_loss_list.append(loss.detach().numpy())
        test_loss_list.append(test_loss.detach().numpy())
        train_epoch_list.append(epoch)
        test_epoch_list.append(epoch)

        Z = model(grid_point).detach().numpy()
        Z = Z[:, 1]
        Z = Z.reshape(xx.shape)

        # 绘制左侧拟合图
        ax1.clear()
        ax1.scatter(class1_points[:, 0], class1_points[:, 1], c="b")
        ax1.scatter(class2_points[:, 0], class2_points[:, 1], c="g")
        ax1.contour(xx, yy, Z, levels=[0.5], colors="red")

        ax2.clear()
        ax2.plot(train_epoch_list, train_loss_list, "r", label="train loss")
        ax2.plot(test_epoch_list, test_loss_list, "b", label="test loss")
        ax2.legend()

        plt.pause(0.2)
plt.show()